Skip to main content

pro_core/
docker.rs

1//! Docker integration for Python projects
2//!
3//! Generate Dockerfiles and build images from `[tool.rx.docker]` config:
4//!
5//! ```toml
6//! [tool.rx.docker]
7//! base-image = "python:3.11-slim"
8//! python-version = "3.11"
9//! entrypoint = ["python", "-m", "myapp"]
10//! cmd = ["--help"]
11//! expose = [8000]
12//! env = { APP_ENV = "production" }
13//! copy = ["src/", "config/"]
14//! workdir = "/app"
15//! user = "appuser"
16//! labels = { maintainer = "dev@example.com" }
17//! apt-packages = ["curl", "gcc"]
18//! build-args = { PIP_NO_CACHE_DIR = "1" }
19//! multi-stage = true
20//! ```
21
22use std::collections::HashMap;
23use std::path::Path;
24
25use crate::pep::PyProject;
26use crate::{Error, Result};
27
28/// Docker configuration from pyproject.toml
29#[derive(Debug, Clone)]
30pub struct DockerConfig {
31    /// Base image (default: python:3.11-slim)
32    pub base_image: String,
33    /// Python version for base image
34    pub python_version: String,
35    /// Working directory in container (default: /app)
36    pub workdir: String,
37    /// Entrypoint command
38    pub entrypoint: Option<Vec<String>>,
39    /// Default command
40    pub cmd: Option<Vec<String>>,
41    /// Ports to expose
42    pub expose: Vec<u16>,
43    /// Environment variables
44    pub env: HashMap<String, String>,
45    /// Additional files/directories to copy
46    pub copy: Vec<String>,
47    /// User to run as
48    pub user: Option<String>,
49    /// Image labels
50    pub labels: HashMap<String, String>,
51    /// APT packages to install
52    pub apt_packages: Vec<String>,
53    /// Build arguments
54    pub build_args: HashMap<String, String>,
55    /// Use multi-stage build
56    pub multi_stage: bool,
57    /// Install dev dependencies
58    pub dev_deps: bool,
59    /// Custom Dockerfile commands (inserted before COPY)
60    pub pre_copy: Vec<String>,
61    /// Custom Dockerfile commands (inserted after COPY)
62    pub post_copy: Vec<String>,
63}
64
65impl Default for DockerConfig {
66    fn default() -> Self {
67        Self {
68            base_image: "python:3.11-slim".to_string(),
69            python_version: "3.11".to_string(),
70            workdir: "/app".to_string(),
71            entrypoint: None,
72            cmd: None,
73            expose: Vec::new(),
74            env: HashMap::new(),
75            copy: Vec::new(),
76            user: None,
77            labels: HashMap::new(),
78            apt_packages: Vec::new(),
79            build_args: HashMap::new(),
80            multi_stage: true,
81            dev_deps: false,
82            pre_copy: Vec::new(),
83            post_copy: Vec::new(),
84        }
85    }
86}
87
88impl DockerConfig {
89    /// Load Docker config from pyproject.toml
90    pub fn load(project_dir: &Path) -> Result<Self> {
91        let pyproject = PyProject::load(project_dir)?;
92
93        let mut config = Self::default();
94
95        let rx_config = match pyproject.tool.get("rx") {
96            Some(c) => c,
97            None => return Ok(config),
98        };
99
100        let docker_config = match rx_config.get("docker") {
101            Some(c) => c,
102            None => return Ok(config),
103        };
104
105        if let Some(v) = docker_config.get("base-image").and_then(|v| v.as_str()) {
106            config.base_image = v.to_string();
107        }
108
109        if let Some(v) = docker_config.get("python-version").and_then(|v| v.as_str()) {
110            config.python_version = v.to_string();
111            // Update base image if not explicitly set
112            if docker_config.get("base-image").is_none() {
113                config.base_image = format!("python:{}-slim", v);
114            }
115        }
116
117        if let Some(v) = docker_config.get("workdir").and_then(|v| v.as_str()) {
118            config.workdir = v.to_string();
119        }
120
121        if let Some(arr) = docker_config.get("entrypoint").and_then(|v| v.as_array()) {
122            config.entrypoint = Some(
123                arr.iter()
124                    .filter_map(|v| v.as_str().map(String::from))
125                    .collect(),
126            );
127        }
128
129        if let Some(arr) = docker_config.get("cmd").and_then(|v| v.as_array()) {
130            config.cmd = Some(
131                arr.iter()
132                    .filter_map(|v| v.as_str().map(String::from))
133                    .collect(),
134            );
135        }
136
137        if let Some(arr) = docker_config.get("expose").and_then(|v| v.as_array()) {
138            config.expose = arr
139                .iter()
140                .filter_map(|v| v.as_integer().map(|i| i as u16))
141                .collect();
142        }
143
144        if let Some(table) = docker_config.get("env").and_then(|v| v.as_table()) {
145            for (k, v) in table {
146                if let Some(s) = v.as_str() {
147                    config.env.insert(k.clone(), s.to_string());
148                }
149            }
150        }
151
152        if let Some(arr) = docker_config.get("copy").and_then(|v| v.as_array()) {
153            config.copy = arr
154                .iter()
155                .filter_map(|v| v.as_str().map(String::from))
156                .collect();
157        }
158
159        if let Some(v) = docker_config.get("user").and_then(|v| v.as_str()) {
160            config.user = Some(v.to_string());
161        }
162
163        if let Some(table) = docker_config.get("labels").and_then(|v| v.as_table()) {
164            for (k, v) in table {
165                if let Some(s) = v.as_str() {
166                    config.labels.insert(k.clone(), s.to_string());
167                }
168            }
169        }
170
171        if let Some(arr) = docker_config.get("apt-packages").and_then(|v| v.as_array()) {
172            config.apt_packages = arr
173                .iter()
174                .filter_map(|v| v.as_str().map(String::from))
175                .collect();
176        }
177
178        if let Some(table) = docker_config.get("build-args").and_then(|v| v.as_table()) {
179            for (k, v) in table {
180                if let Some(s) = v.as_str() {
181                    config.build_args.insert(k.clone(), s.to_string());
182                }
183            }
184        }
185
186        if let Some(v) = docker_config.get("multi-stage").and_then(|v| v.as_bool()) {
187            config.multi_stage = v;
188        }
189
190        if let Some(v) = docker_config.get("dev-deps").and_then(|v| v.as_bool()) {
191            config.dev_deps = v;
192        }
193
194        if let Some(arr) = docker_config.get("pre-copy").and_then(|v| v.as_array()) {
195            config.pre_copy = arr
196                .iter()
197                .filter_map(|v| v.as_str().map(String::from))
198                .collect();
199        }
200
201        if let Some(arr) = docker_config.get("post-copy").and_then(|v| v.as_array()) {
202            config.post_copy = arr
203                .iter()
204                .filter_map(|v| v.as_str().map(String::from))
205                .collect();
206        }
207
208        Ok(config)
209    }
210}
211
212/// Dockerfile generator
213pub struct DockerfileGenerator {
214    config: DockerConfig,
215    project_name: String,
216}
217
218impl DockerfileGenerator {
219    /// Create a new generator
220    pub fn new(config: DockerConfig, project_name: String) -> Self {
221        Self {
222            config,
223            project_name,
224        }
225    }
226
227    /// Load from project directory
228    pub fn from_project(project_dir: &Path) -> Result<Self> {
229        let config = DockerConfig::load(project_dir)?;
230        let pyproject = PyProject::load(project_dir)?;
231        let project_name = pyproject
232            .name()
233            .unwrap_or("app")
234            .to_string()
235            .replace('-', "_");
236
237        Ok(Self::new(config, project_name))
238    }
239
240    /// Generate Dockerfile content
241    pub fn generate(&self) -> String {
242        if self.config.multi_stage {
243            self.generate_multi_stage()
244        } else {
245            self.generate_single_stage()
246        }
247    }
248
249    /// Generate a multi-stage Dockerfile (smaller final image)
250    fn generate_multi_stage(&self) -> String {
251        let mut lines = Vec::new();
252
253        // Build stage
254        lines.push(format!(
255            "# Build stage\nFROM {} AS builder",
256            self.config.base_image
257        ));
258        lines.push(String::new());
259
260        // Build args
261        for (key, value) in &self.config.build_args {
262            lines.push(format!("ARG {}={}", key, value));
263        }
264        if !self.config.build_args.is_empty() {
265            lines.push(String::new());
266        }
267
268        // Install build dependencies
269        lines.push("# Install build dependencies".to_string());
270        let mut apt_deps = vec!["build-essential"];
271        apt_deps.extend(self.config.apt_packages.iter().map(|s| s.as_str()));
272
273        lines.push(format!(
274            "RUN apt-get update && apt-get install -y --no-install-recommends {} && rm -rf /var/lib/apt/lists/*",
275            apt_deps.join(" ")
276        ));
277        lines.push(String::new());
278
279        // Set workdir
280        lines.push(format!("WORKDIR {}", self.config.workdir));
281        lines.push(String::new());
282
283        // Copy and install dependencies first (for layer caching)
284        lines.push("# Install dependencies".to_string());
285        lines.push("COPY pyproject.toml ./".to_string());
286        lines.push("COPY rx.lock* ./".to_string());
287        lines.push(String::new());
288
289        // Create virtual environment and install deps
290        lines.push("RUN python -m venv /opt/venv".to_string());
291        lines.push("ENV PATH=\"/opt/venv/bin:$PATH\"".to_string());
292        lines.push(String::new());
293
294        // Install using rx if available, otherwise pip
295        lines.push("RUN pip install --no-cache-dir --upgrade pip && \\".to_string());
296        lines.push(
297            "    if [ -f rx.lock ]; then pip install --no-cache-dir -r <(python -c \"import tomllib; f=open('rx.lock','rb'); d=tomllib.load(f); print('\\\\n'.join(f\\\"{p}=={d['packages'][p]['version']}\\\" for p in d['packages']))\") 2>/dev/null || pip install .; else pip install .; fi"
298                .to_string(),
299        );
300        lines.push(String::new());
301
302        // Custom pre-copy commands
303        for cmd in &self.config.pre_copy {
304            lines.push(format!("RUN {}", cmd));
305        }
306        if !self.config.pre_copy.is_empty() {
307            lines.push(String::new());
308        }
309
310        // Copy source code
311        lines.push("# Copy source code".to_string());
312        lines.push("COPY . .".to_string());
313        lines.push(String::new());
314
315        // Install the package itself
316        lines.push("RUN pip install --no-cache-dir .".to_string());
317        lines.push(String::new());
318
319        // Runtime stage
320        lines.push(format!(
321            "# Runtime stage\nFROM {} AS runtime",
322            self.config.base_image
323        ));
324        lines.push(String::new());
325
326        // Labels
327        for (key, value) in &self.config.labels {
328            lines.push(format!("LABEL {}=\"{}\"", key, value));
329        }
330        if !self.config.labels.is_empty() {
331            lines.push(String::new());
332        }
333
334        // Install runtime apt packages only
335        if !self.config.apt_packages.is_empty() {
336            let runtime_pkgs: Vec<_> = self
337                .config
338                .apt_packages
339                .iter()
340                .filter(|p| !["build-essential", "gcc", "g++", "make"].contains(&p.as_str()))
341                .collect();
342
343            if !runtime_pkgs.is_empty() {
344                lines.push(format!(
345                    "RUN apt-get update && apt-get install -y --no-install-recommends {} && rm -rf /var/lib/apt/lists/*",
346                    runtime_pkgs.iter().map(|s| s.as_str()).collect::<Vec<_>>().join(" ")
347                ));
348                lines.push(String::new());
349            }
350        }
351
352        // Set workdir
353        lines.push(format!("WORKDIR {}", self.config.workdir));
354        lines.push(String::new());
355
356        // Copy venv from builder
357        lines.push("# Copy virtual environment from builder".to_string());
358        lines.push("COPY --from=builder /opt/venv /opt/venv".to_string());
359        lines.push("ENV PATH=\"/opt/venv/bin:$PATH\"".to_string());
360        lines.push(String::new());
361
362        // Copy additional files
363        if !self.config.copy.is_empty() {
364            lines.push("# Copy additional files".to_string());
365            for path in &self.config.copy {
366                lines.push(format!("COPY {} {}/", path, self.config.workdir));
367            }
368            lines.push(String::new());
369        }
370
371        // Custom post-copy commands
372        for cmd in &self.config.post_copy {
373            lines.push(format!("RUN {}", cmd));
374        }
375        if !self.config.post_copy.is_empty() {
376            lines.push(String::new());
377        }
378
379        // Environment variables
380        for (key, value) in &self.config.env {
381            lines.push(format!("ENV {}=\"{}\"", key, value));
382        }
383        if !self.config.env.is_empty() {
384            lines.push(String::new());
385        }
386
387        // Expose ports
388        for port in &self.config.expose {
389            lines.push(format!("EXPOSE {}", port));
390        }
391        if !self.config.expose.is_empty() {
392            lines.push(String::new());
393        }
394
395        // User
396        if let Some(ref user) = self.config.user {
397            lines.push(format!("RUN useradd -m -s /bin/bash {}", user));
398            lines.push(format!("USER {}", user));
399            lines.push(String::new());
400        }
401
402        // Entrypoint and CMD
403        if let Some(ref entrypoint) = self.config.entrypoint {
404            let ep_str = entrypoint
405                .iter()
406                .map(|s| format!("\"{}\"", s))
407                .collect::<Vec<_>>()
408                .join(", ");
409            lines.push(format!("ENTRYPOINT [{}]", ep_str));
410        }
411
412        if let Some(ref cmd) = self.config.cmd {
413            let cmd_str = cmd
414                .iter()
415                .map(|s| format!("\"{}\"", s))
416                .collect::<Vec<_>>()
417                .join(", ");
418            lines.push(format!("CMD [{}]", cmd_str));
419        } else if self.config.entrypoint.is_none() {
420            // Default: run the package as a module
421            lines.push(format!(
422                "CMD [\"python\", \"-m\", \"{}\"]",
423                self.project_name
424            ));
425        }
426
427        lines.join("\n")
428    }
429
430    /// Generate a single-stage Dockerfile
431    fn generate_single_stage(&self) -> String {
432        let mut lines = Vec::new();
433
434        lines.push(format!("FROM {}", self.config.base_image));
435        lines.push(String::new());
436
437        // Build args
438        for (key, value) in &self.config.build_args {
439            lines.push(format!("ARG {}={}", key, value));
440        }
441        if !self.config.build_args.is_empty() {
442            lines.push(String::new());
443        }
444
445        // Labels
446        for (key, value) in &self.config.labels {
447            lines.push(format!("LABEL {}=\"{}\"", key, value));
448        }
449        if !self.config.labels.is_empty() {
450            lines.push(String::new());
451        }
452
453        // Install apt packages
454        if !self.config.apt_packages.is_empty() {
455            lines.push(format!(
456                "RUN apt-get update && apt-get install -y --no-install-recommends {} && rm -rf /var/lib/apt/lists/*",
457                self.config.apt_packages.join(" ")
458            ));
459            lines.push(String::new());
460        }
461
462        // Set workdir
463        lines.push(format!("WORKDIR {}", self.config.workdir));
464        lines.push(String::new());
465
466        // Copy requirements first for caching
467        lines.push("# Copy dependency files".to_string());
468        lines.push("COPY pyproject.toml ./".to_string());
469        lines.push("COPY rx.lock* ./".to_string());
470        lines.push(String::new());
471
472        // Install dependencies
473        lines.push("# Install dependencies".to_string());
474        lines.push("RUN pip install --no-cache-dir --upgrade pip".to_string());
475        lines.push(String::new());
476
477        // Custom pre-copy commands
478        for cmd in &self.config.pre_copy {
479            lines.push(format!("RUN {}", cmd));
480        }
481        if !self.config.pre_copy.is_empty() {
482            lines.push(String::new());
483        }
484
485        // Copy source code
486        lines.push("# Copy source code".to_string());
487        lines.push("COPY . .".to_string());
488        lines.push(String::new());
489
490        // Install the package
491        lines.push("RUN pip install --no-cache-dir .".to_string());
492        lines.push(String::new());
493
494        // Copy additional files
495        if !self.config.copy.is_empty() {
496            for path in &self.config.copy {
497                lines.push(format!("COPY {} {}/", path, self.config.workdir));
498            }
499            lines.push(String::new());
500        }
501
502        // Custom post-copy commands
503        for cmd in &self.config.post_copy {
504            lines.push(format!("RUN {}", cmd));
505        }
506        if !self.config.post_copy.is_empty() {
507            lines.push(String::new());
508        }
509
510        // Environment variables
511        for (key, value) in &self.config.env {
512            lines.push(format!("ENV {}=\"{}\"", key, value));
513        }
514        if !self.config.env.is_empty() {
515            lines.push(String::new());
516        }
517
518        // Expose ports
519        for port in &self.config.expose {
520            lines.push(format!("EXPOSE {}", port));
521        }
522        if !self.config.expose.is_empty() {
523            lines.push(String::new());
524        }
525
526        // User
527        if let Some(ref user) = self.config.user {
528            lines.push(format!("RUN useradd -m -s /bin/bash {}", user));
529            lines.push(format!("USER {}", user));
530            lines.push(String::new());
531        }
532
533        // Entrypoint and CMD
534        if let Some(ref entrypoint) = self.config.entrypoint {
535            let ep_str = entrypoint
536                .iter()
537                .map(|s| format!("\"{}\"", s))
538                .collect::<Vec<_>>()
539                .join(", ");
540            lines.push(format!("ENTRYPOINT [{}]", ep_str));
541        }
542
543        if let Some(ref cmd) = self.config.cmd {
544            let cmd_str = cmd
545                .iter()
546                .map(|s| format!("\"{}\"", s))
547                .collect::<Vec<_>>()
548                .join(", ");
549            lines.push(format!("CMD [{}]", cmd_str));
550        } else if self.config.entrypoint.is_none() {
551            lines.push(format!(
552                "CMD [\"python\", \"-m\", \"{}\"]",
553                self.project_name
554            ));
555        }
556
557        lines.join("\n")
558    }
559
560    /// Generate .dockerignore content
561    pub fn generate_dockerignore(&self) -> String {
562        r#"# Python
563__pycache__/
564*.py[cod]
565*$py.class
566*.so
567.Python
568build/
569develop-eggs/
570dist/
571downloads/
572eggs/
573.eggs/
574lib/
575lib64/
576parts/
577sdist/
578var/
579wheels/
580*.egg-info/
581.installed.cfg
582*.egg
583
584# Virtual environments
585.venv/
586venv/
587ENV/
588env/
589
590# IDE
591.idea/
592.vscode/
593*.swp
594*.swo
595
596# Testing
597.tox/
598.nox/
599.coverage
600htmlcov/
601.pytest_cache/
602.mypy_cache/
603
604# Git
605.git/
606.gitignore
607
608# Docker
609Dockerfile
610.dockerignore
611docker-compose*.yml
612
613# Documentation
614docs/
615*.md
616!README.md
617
618# Misc
619*.log
620.DS_Store
621Thumbs.db
622"#
623        .to_string()
624    }
625}
626
627/// Build a Docker image
628pub fn build_image(
629    project_dir: &Path,
630    tag: &str,
631    dockerfile_path: Option<&Path>,
632    build_args: &HashMap<String, String>,
633    no_cache: bool,
634) -> Result<()> {
635    use std::process::Command;
636
637    let mut cmd = Command::new("docker");
638    cmd.arg("build");
639    cmd.arg("-t").arg(tag);
640
641    if let Some(df) = dockerfile_path {
642        cmd.arg("-f").arg(df);
643    }
644
645    for (key, value) in build_args {
646        cmd.arg("--build-arg").arg(format!("{}={}", key, value));
647    }
648
649    if no_cache {
650        cmd.arg("--no-cache");
651    }
652
653    cmd.arg(project_dir);
654    cmd.current_dir(project_dir);
655
656    let status = cmd
657        .status()
658        .map_err(|e| Error::Config(format!("Failed to run docker build: {}", e)))?;
659
660    if !status.success() {
661        return Err(Error::Config("Docker build failed".to_string()));
662    }
663
664    Ok(())
665}
666
667#[cfg(test)]
668mod tests {
669    use super::*;
670
671    #[test]
672    fn test_default_config() {
673        let config = DockerConfig::default();
674        assert_eq!(config.base_image, "python:3.11-slim");
675        assert_eq!(config.workdir, "/app");
676        assert!(config.multi_stage);
677    }
678
679    #[test]
680    fn test_generate_dockerfile() {
681        let config = DockerConfig {
682            base_image: "python:3.11-slim".to_string(),
683            python_version: "3.11".to_string(),
684            workdir: "/app".to_string(),
685            expose: vec![8000],
686            env: [("APP_ENV".to_string(), "production".to_string())]
687                .into_iter()
688                .collect(),
689            multi_stage: false,
690            ..Default::default()
691        };
692
693        let generator = DockerfileGenerator::new(config, "myapp".to_string());
694        let dockerfile = generator.generate();
695
696        assert!(dockerfile.contains("FROM python:3.11-slim"));
697        assert!(dockerfile.contains("WORKDIR /app"));
698        assert!(dockerfile.contains("EXPOSE 8000"));
699        assert!(dockerfile.contains("ENV APP_ENV=\"production\""));
700    }
701
702    #[test]
703    fn test_multi_stage_dockerfile() {
704        let config = DockerConfig {
705            multi_stage: true,
706            ..Default::default()
707        };
708
709        let generator = DockerfileGenerator::new(config, "myapp".to_string());
710        let dockerfile = generator.generate();
711
712        assert!(dockerfile.contains("AS builder"));
713        assert!(dockerfile.contains("AS runtime"));
714        assert!(dockerfile.contains("COPY --from=builder"));
715    }
716}